import cv2
import os
import numpy as np
from PIL import Image, ImageDraw, ImageFont


def add_chinese(img, text, position, text_color, text_size):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # create a Draw object
    draw = ImageDraw.Draw(img)
    # font format
    path = os.path.dirname(os.path.abspath(__file__))
    font_path = os.path.join(path, "font", "simsun.ttc")
    font_style = ImageFont.truetype(font_path, int(text_size))
    # draw chinese
    draw.text(position, text, text_color, font=font_style)
    # convert to format in opencv
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def label2chinese(image, x0, y0, x1, y1, name):
    """
    :param image: image to be drawn with chinese
    :param x0: left coordinate
    :param y0: top coordinate
    :param x1: right coordinate
    :param y1: bottom coordinate
    :param name: chinese text
    :return:
    """
    width = x1 - x0
    # txt_color = (0, 0, 255) if numpy.mean(_color) > 0.5 else (0, 0, 255)
    font = cv2.FONT_HERSHEY_SIMPLEX
    txt_size = cv2.getTextSize(name, font, 0.8, 1)[0]
    txt_bk_color = (0, 215, 255)
    cv2.rectangle(
        image,
        (x0 + int(width / 2) - int(0.25 * txt_size[0]), y0 - int(1.3 * txt_size[1])),
        (x0 + int(width / 2) + int(0.25 * txt_size[0]), y0),
        txt_bk_color,
        -1)
    try:
        image = add_chinese(image,
                            name,
                            (x0 + int(width / 2) - int(0.25 * txt_size[0]), y0 - int(1.3 * txt_size[1])),
                            (0, 0, 0),
                            20)
    except Exception as e:
        print("Exception occurred in draw yolo while adding chinese : ", e)
    return image


def draw_chinese(img, text, positive, font_size=20, font_color=(0, 0, 0)):
    # cv2和PIL中颜色的hex码的储存顺序不同
    cv2img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    pilimg = Image.fromarray(cv2img)
    # PIL图片上打印汉字
    draw = ImageDraw.Draw(pilimg)  # 图片上打印
    # 参数1：字体文件路径，参数2：字体大小
    font = ImageFont.truetype("data/ChineseFont/FangZhengKaiTiJianTi-1.ttf", font_size, encoding="utf-8")
    # 参数1：打印坐标，参数2：文本，参数3：字体颜色，参数4：字体格式
    draw.text(positive, text, font_color, font=font)
    cv2charimg = cv2.cvtColor(np.array(pilimg), cv2.COLOR_RGB2BGR)  # PIL图片转cv2 图片
    return cv2charimg


def vis(img, boxes, scores, cls_ids, class_names=None):
    for i in range(len(boxes)):
        box = boxes[i]
        cls_id = int(cls_ids[i])
        score = scores[i]
        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])

        color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
        cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)

        text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)
        txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
        font = cv2.FONT_HERSHEY_SIMPLEX

        txt_size = cv2.getTextSize(text, font, 1.0, 1)[0]
        txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()
        cv2.rectangle(
            img,
            (x0, y0),
            (round(x0 + 0.5 * txt_size[0]), y0 + round(0.5 * txt_size[1])),
            txt_bk_color,
            -1)
        cv2.putText(img, text, (x0, round(y0 + 0.5 * txt_size[1])), font, 0.5, txt_color, thickness=1)
    return img


def vis_track(img, boxes):
    for i in range(len(boxes)):
        box = boxes[i]

        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])

        id = box[4]
        color_ = _COLORS[id % _COLORS.shape[0]]
        color = (color_ * 255).astype(np.uint8).tolist()
        text = '%d' % (id)
        if (len(box) > 5):
            text = box[5]
        if (len(box) > 6):
            prob = int(100 * box[6])
            if (prob < 70):
                continue
            else:
                text = text + " " + str(prob) + "%"
        txt_color = (255, 255, 255)
        font = cv2.FONT_HERSHEY_SIMPLEX

        txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
        cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)

        txt_bk_color = (color_ * 255 * 0.7).astype(np.uint8).tolist()
        cv2.rectangle(
            img,
            (x0, y0 + 1),
            (x0 + txt_size[0] + 1, y0 + int(1.5 * txt_size[1])),
            txt_bk_color,
            -1
        )
        cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)
    return img


_COLORS = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.300, 0.300, 0.300,
        0.600, 0.600, 0.600,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.749, 0.749, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 1.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.286, 0.286, 0.286,
        0.429, 0.429, 0.429,
        0.571, 0.571, 0.571,
        0.714, 0.714, 0.714,
        0.857, 0.857, 0.857,
        0.000, 0.447, 0.741,
        0.314, 0.717, 0.741,
        0.50, 0.5, 0
    ]
).astype(np.float32).reshape(-1, 3)
